# 本示例演示如何使用kan拟合函数，即回归

import torch
from   kan import KAN
from   kan import create_dataset

#! 1. 新建一个kan模型，
"""创建一个kan模型(继承nn.Module类)
@param width = [2, 5, 1]  kan模型的mlp模型结构
@param grid  = 5          5个网格区间   
@param k     = 3          样条函数的阶数
@param seed  = 0          随机种子
"""
model = KAN(width = [2, 5, 1], grid = 5, k = 3, seed = 0)


#! 2. 创建数据集
# 本示例用于拟合: f(x, y) = exp(sin(pi*x)+ y^2)
func    = lambda x: torch.exp(torch.sin(torch.pi * x[:,[0]]) + x[:,[1]]**2)
dataset = create_dataset(func, n_var=2, train_num=1000, test_num=1000)
    # dataset['train_input'] = train_input.to(device)
    # dataset['test_input'] = test_input.to(device)

    # dataset['train_label'] = train_label.to(device)
    # dataset['test_label'] = test_label.to(device)
print(dataset["train_input"].shape, dataset["train_label"].shape,
      dataset["test_input"].shape, dataset["test_label"].shape)

#! 2. 可视化模型的初始状态
# model(dataset["train_input"])  # 推理一次模型是的模型拥有激活值
# model.plot(beta=100)

#! 3. 训练模型
"""模型训练
@param dataset       模型训练数据集
@param opt           训练优化算法, "LBFGS" or "Adam"
@param steps         训练步数
@param lamb          全局惩罚强度
@param lamb_entropy  交叉熵惩罚强度
"""
model.train(dataset, opt="LBFGS", steps=20, lamb=0.01, lamb_entropy=10.0)
# model.plot()

# 模型剪枝并按照原始的结构显示
# model.prune()
# model.plot(mask=True)

# 模型剪枝并按精简后的结构显示
model(dataset["train_input"])
model = model.prune()         # 获取键之后的模型
# model.plot(mask=True)

#! 4. 模型剪枝后再训练
model.train(dataset, opt="LBFGS", steps=50)
model.plot()

#! 5. 手动指定激活函数
mode = "auto" # "manual"

if mode == "manual":         # 手动指定激活函数形式
    model.fix_symbolic(0,0,0,'sin')
    model.fix_symbolic(0,1,0,'x^2')
    model.fix_symbolic(1,0,0,'exp')
elif mode == "auto":         # 自动模式, 在给定列表中选择
    lib = ['x','x^2','x^3','x^4','exp','log','sqrt','tanh','sin','abs']
    model.auto_symbolic(lib=lib)
# 配置激活函数后继续训练
model.train(dataset, opt="LBFGS", steps=50)


#! 6. 获取符号公式
print(model.symbolic_formula())